# __init__.py
import os
import sys
from pathlib import Path

# 使用传统方式，但更加健壮
PROJECT_ROOT_STR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
if PROJECT_ROOT_STR not in sys.path:
    sys.path.insert(0, PROJECT_ROOT_STR)
    
from config.path_config import *
"""PPQ量化相关工具"""
import json
import tensorrt as trt
from typing import Dict, Any
from core.tensorrt_core import INT8Strategy

class PPQJsonLoader:
    """PPQ JSON文件加载器"""
    
    @staticmethod
    def load_json(json_path: str) -> Dict[str, Any]:
        """加载JSON文件"""
        try:
            with open(json_path, 'r') as f:
                return json.load(f)
        except (FileNotFoundError, json.JSONDecodeError) as e:
            raise RuntimeError(f"Failed to load PPQ JSON file {json_path}: {e}")

class PPQDynamicRangeSetter:
    """PPQ动态范围设置器"""
    
    def __init__(self, ppq_json_path: str):
        self.quant_params = PPQJsonLoader.load_json(ppq_json_path)
        self.act_quant = self.quant_params.get("act_quant_info", {})
    
    def set_dynamic_ranges(self, network: trt.INetworkDefinition) -> None:
        """为网络设置动态范围"""
        # 设置输入张量动态范围
        self._set_input_ranges(network)
        
        # 设置输出张量动态范围
        self._set_output_ranges(network)
    
    def _set_input_ranges(self, network: trt.INetworkDefinition) -> None:
        """设置输入张量动态范围"""
        for i in range(network.num_inputs):
            input_tensor = network.get_input(i)
            if input_tensor.name in self.act_quant:
                self._apply_dynamic_range(input_tensor, self.act_quant[input_tensor.name])
    
    def _set_output_ranges(self, network: trt.INetworkDefinition) -> None:
        """设置输出张量动态范围"""
        for i in range(network.num_layers):
            layer = network.get_layer(i)
            for output_index in range(layer.num_outputs):
                tensor = layer.get_output(output_index)
                if tensor.name in self.act_quant:
                    self._apply_dynamic_range(tensor, self.act_quant[tensor.name])
                else:
                    print(f"\033[1;32mMissing quantization info for: {tensor.name}\033[0m")
    
    def _apply_dynamic_range(self, tensor: trt.ITensor, value: float) -> None:
        """应用动态范围到张量"""
        tensor_max = abs(value)
        tensor_min = -abs(value)
        tensor.set_dynamic_range(tensor_min, tensor_max)

class PPQINT8Strategy(INT8Strategy):
    """PPQ专用INT8量化策略"""
    
    def __init__(self, ppq_json_path: str):
        super().__init__(calibrator=None)  # PPQ不需要校准器
        self.ppq_json_path = ppq_json_path
        self.range_setter = PPQDynamicRangeSetter(ppq_json_path)
    
    def apply_quantization(self, config: trt.IBuilderConfig, network: trt.INetworkDefinition) -> None:
        super().apply_quantization(config, network)
        self.range_setter.set_dynamic_ranges(network)